package com.ob.api.limit;

import cn.hutool.core.annotation.AnnotationUtil;
import com.ob.api.exception.ServerUnavailableException;
import com.ob.core.constants.enums.SysResponseCode;
import com.ob.core.dto.R;
import com.ob.core.util.JsonUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * @author xzzz
 */
public class RateLimiterInterceptor implements HandlerInterceptor {
    private static final Logger log = LoggerFactory.getLogger(RateLimiterInterceptor.class);

    @Autowired
    private RateLimiterIp rateLimiterIp;

    @Autowired
    private RateLimiterApi rateLimiterApi;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {

        HandlerMethod handlerMethod = (HandlerMethod) handler;
        RateLimiter rateLimiter = AnnotationUtil.getAnnotation(handlerMethod.getMethod(), RateLimiter.class);
        String limiterType = checkLimiterType(rateLimiter);

        if (LimiterType.IP.name().equals(limiterType)) {
            if (rateLimiterIp.limiting(request,rateLimiter)) {
                responseWrap(response);
                return false;
            }
        }

        // TODO
        if (LimiterType.API.name().equals(limiterType)) {
            log.info("API限流");
        }

        return true;
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) {
    }

    /**
     * 清除Token ThreadLocal 上下文
     */
    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) {
    }

    private String checkLimiterType (RateLimiter rateLimiter) {
        if (rateLimiter == null) {
            return null;
        }

        if (LimiterType.IP.name().equals(rateLimiter.type())) {
            return LimiterType.IP.name();
        }

        if (LimiterType.API.name().equals(rateLimiter.type())) {
            return LimiterType.API.name();
        }

        return null;
    }

    private void responseWrap (HttpServletResponse response) throws IOException {
        SysResponseCode resp503 = ServerUnavailableException.resp;
        response.setContentType(MediaType.APPLICATION_JSON_VALUE + ";charset=UTF-8");
        response.setStatus(HttpStatus.SERVICE_UNAVAILABLE.value());
        PrintWriter out = response.getWriter();
        out.print(JsonUtil.obj2Json(R.fault(resp503.code(),resp503.msg())));
    }

}
